
#include <map>
#include <ruby.h>
#include <udt.h>

static VALUE eUDTSocket;

static std::map<VALUE, UDTSOCKET> fds;

// helper functions

static int udt_get_int(UDTSOCKET fd, UDTOpt opt)
{
	int ret, size = sizeof(ret);
	if(UDT::getsockopt(fd, 0, opt, &ret, &size) != 0)
		rb_raise(eUDTSocket, "UDT::getsockopt() failed: %s", UDT::getlasterror().getErrorMessage());
	return ret;
}

static void udt_set_int(UDTSOCKET fd, UDTOpt opt, int val)
{
	if(UDT::setsockopt(fd, 0, opt, &val, sizeof(val)) != 0)
		rb_raise(eUDTSocket, "UDT::setsockopt() failed: %s", UDT::getlasterror().getErrorMessage());
}

static int64_t udt_get_int64(UDTSOCKET fd, UDTOpt opt)
{
	int64_t ret;
	int size = sizeof(ret);
	if(UDT::getsockopt(fd, 0, opt, &ret, &size) != 0)
		rb_raise(eUDTSocket, "UDT::getsockopt() failed: %s", UDT::getlasterror().getErrorMessage());
	return ret;
}

static void udt_set_int64(UDTSOCKET fd, UDTOpt opt, int64_t val)
{
	if(UDT::setsockopt(fd, 0, opt, &val, sizeof(val)) != 0)
		rb_raise(eUDTSocket, "UDT::setsockopt() failed: %s", UDT::getlasterror().getErrorMessage());
}

static bool udt_get_bool(UDTSOCKET fd, UDTOpt opt)
{
	bool ret;
	int size = sizeof(ret);
	if(UDT::getsockopt(fd, 0, opt, &ret, &size) != 0)
		rb_raise(eUDTSocket, "UDT::getsockopt() failed: %s", UDT::getlasterror().getErrorMessage());
	return ret;
}

static void udt_set_bool(UDTSOCKET fd, UDTOpt opt, bool val)
{
	if(UDT::setsockopt(fd, 0, opt, &val, sizeof(val)) != 0)
		rb_raise(eUDTSocket, "UDT::setsockopt() failed: %s", UDT::getlasterror().getErrorMessage());
}

// udt bindings

static VALUE udt_s_startup(VALUE klass)
{
	if(UDT::startup() != 0)
	{
		if(UDT::getlasterror().getErrorCode() != CUDTException::SUCCESS)
			rb_raise(eUDTSocket, "UDT::startup() failed: %s", UDT::getlasterror().getErrorMessage());
	}
	return INT2FIX(0);
}

static VALUE udt_s_cleanup(VALUE klass)
{
	if(UDT::cleanup() != 0)
	{
		if(UDT::getlasterror().getErrorCode() != CUDTException::SUCCESS)
			rb_raise(eUDTSocket, "UDT::cleanup() failed: %s", UDT::getlasterror().getErrorMessage());
	}
	return INT2FIX(0);
}

static VALUE udt_socket(VALUE klass, VALUE domain, VALUE type)
{
	VALUE tmp;
	char* ptr;
	int d, t;
	
	tmp  = rb_check_string_type(domain);
	if(!NIL_P(tmp))
	{
		domain = tmp;
		rb_check_safe_obj(domain);
		ptr = RSTRING(domain)->as.heap.ptr;
		if(strcmp(ptr, "AF_INET") == 0)
			d = AF_INET;
		else if(strcmp(ptr, "AF_INET6") == 0)
			d = AF_INET6;
		else
			rb_raise(eUDTSocket, "unknown socket domain %s", ptr);
	}
	else
		d = NUM2INT(domain);
	tmp  = rb_check_string_type(type);
	if(!NIL_P(tmp))
	{
		type = tmp;
		rb_check_safe_obj(type);
		ptr = RSTRING(type)->as.heap.ptr;
		if(strcmp(ptr, "SOCK_STREAM") == 0)
			t = SOCK_STREAM;
		else if(strcmp(ptr, "SOCK_DGRAM") == 0)
			t = SOCK_DGRAM;
		else
			rb_raise(eUDTSocket, "unknown socket type %s", ptr);
	}
	else
		t = NUM2INT(type);
	UDTSOCKET fd = UDT::socket(d, t, 0);
	if(fd == UDT::INVALID_SOCK)
		rb_raise(eUDTSocket, "UDT::socket() failed: %s", UDT::getlasterror().getErrorMessage());
	fds[klass] = fd;
	return klass;
}

static VALUE udt_close(VALUE klass)
{
	UDTSOCKET fd = fds[klass];
	if(UDT::close(fd) != 0)
		rb_raise(eUDTSocket, "UDT::close() failed: %s", UDT::getlasterror().getErrorMessage());
	return INT2FIX(0);
}

static VALUE udt_connect(VALUE klass, VALUE addr)
{
	StringValue(addr);
    addr = rb_str_new4(addr);
	UDTSOCKET fd = fds[klass];
	if(UDT::connect(fd, (struct sockaddr*)RSTRING(addr)->as.heap.ptr, RSTRING(addr)->as.heap.len) != 0)
		rb_raise(eUDTSocket, "UDT::connect() failed: %s", UDT::getlasterror().getErrorMessage());
	return INT2FIX(0);
}

static VALUE udt_accept(VALUE klass)
{
	struct sockaddr_storage addr;
	int len;
	int ret = UDT::accept(fds[klass], (struct sockaddr*)(&addr), &len);
	if(ret == UDT::INVALID_SOCK)
	{
		if(UDT::getlasterror().getErrorCode() == CUDTException::EASYNCRCV)
			return Qnil;
		else
			rb_raise(eUDTSocket, "UDT::accept() failed: %s", UDT::getlasterror().getErrorMessage());
	}
	// how to allocate an object for this type???
	VALUE obj = rb_obj_alloc(klass);
	fds[obj] = ret;
	return obj;
}

static VALUE udt_getsockname(VALUE klass)
{
	int len;
	char buff[sizeof(struct sockaddr_storage)];
	if(UDT::getsockname(fds[klass], (struct sockaddr*)buff, &len) != 0)
		rb_raise(eUDTSocket, "UDT::getsockname() failed: %s", UDT::getlasterror().getErrorMessage());
	return rb_str_new(buff, len);
}

static VALUE udt_getpeername(VALUE klass)
{
	int len;
	char buff[sizeof(struct sockaddr_storage)];
	if(UDT::getpeername(fds[klass], (struct sockaddr*)buff, &len) != 0)
		rb_raise(eUDTSocket, "UDT::getpeername() failed: %s", UDT::getlasterror().getErrorMessage());
	return rb_str_new(buff, len);
}

static VALUE udt_bind(VALUE klass, VALUE addr)
{
	StringValue(addr);
    if(UDT::bind(fds[klass], (struct sockaddr*)RSTRING(addr)->as.heap.ptr, RSTRING(addr)->as.heap.len) != 0)
		rb_raise(eUDTSocket, "UDT::bind() failed: %s", UDT::getlasterror().getErrorMessage());
	return INT2FIX(0);
}

static VALUE udt_attach(VALUE klass, VALUE udpsock)
{
	// how to find UDPSocket and its fd???
	return INT2FIX(0);
}

static VALUE udt_listen(VALUE klass, VALUE backlog)
{
	if(TYPE(backlog) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	if(UDT::listen(fds[klass], NUM2INT(backlog)) != 0)
		rb_raise(eUDTSocket, "UDT::listen() failed: %s", UDT::getlasterror().getErrorMessage());
	return INT2FIX(0);	
}

static VALUE udt_send(VALUE klass, VALUE mesg)
{
	StringValue(mesg);
	int n = UDT::send(fds[klass], RSTRING(mesg)->as.heap.ptr, RSTRING(mesg)->as.heap.len, 0);
	return INT2FIX(n);
}

static VALUE udt_recv(VALUE klass, VALUE len)
{
	if(TYPE(len) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	VALUE str = rb_tainted_str_new(0, NUM2INT(len));
	int ret = UDT::recv(fds[klass], RSTRING(str)->as.heap.ptr, NUM2INT(len), 0);
	if(ret == UDT::ERROR)
	{
		UDT::ERRORINFO& e = UDT::getlasterror();
		if(e.getErrorCode() == CUDTException::EASYNCRCV)
			ret = 0;
		else
			rb_raise(eUDTSocket, "UDT::recv() failed: %s", e.getErrorMessage());
	}
	return INT2FIX(ret);
}

static VALUE udt_select(int argc, VALUE* argv, VALUE klass)
{
	return INT2FIX(0);
}

// socket options as properties

static VALUE udt_mss_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDT_MSS));
}

static VALUE udt_mss_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDT_MSS, FIX2INT(val));
	return Qnil;
}

static VALUE udt_sndsyn_get(VALUE klass)
{
	return udt_get_bool(fds[klass], UDT_SNDSYN) ? Qtrue : Qfalse;
}

static VALUE udt_sndsyn_set(VALUE klass, VALUE val)
{
	if((TYPE(val) != T_TRUE) && (TYPE(val) != T_FALSE))
		rb_raise(eUDTSocket, "argument error");
	udt_set_bool(fds[klass], UDT_SNDSYN, (val == Qtrue));
	return Qnil;
}

static VALUE udt_rcvsyn_get(VALUE klass)
{
	return udt_get_bool(fds[klass], UDT_RCVSYN) ? Qtrue : Qfalse;
}

static VALUE udt_rcvsyn_set(VALUE klass, VALUE val)
{
	if((TYPE(val) != T_TRUE) && (TYPE(val) != T_FALSE))
		rb_raise(eUDTSocket, "argument error");
	udt_set_bool(fds[klass], UDT_RCVSYN, (val == Qtrue));
	return Qnil;
}

static VALUE udt_fc_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDT_FC));
}

static VALUE udt_fc_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDT_FC, FIX2INT(val));
	return Qnil;
}

static VALUE udt_sndbuf_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDT_SNDBUF));
}

static VALUE udt_sndbuf_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDT_SNDBUF, FIX2INT(val));
	return Qnil;
}

static VALUE udt_rcvbuf_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDT_RCVBUF));
}

static VALUE udt_rcvbuf_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDT_RCVBUF, FIX2INT(val));
	return Qnil;
}

static VALUE udp_sndbuf_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDP_SNDBUF));
}

static VALUE udp_sndbuf_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDP_SNDBUF, FIX2INT(val));
	return Qnil;
}

static VALUE udp_rcvbuf_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDP_RCVBUF));
}

static VALUE udp_rcvbuf_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDP_RCVBUF, FIX2INT(val));
	return Qnil;
}

static VALUE udt_rendezvous_get(VALUE klass)
{
	return udt_get_bool(fds[klass], UDT_RENDEZVOUS) ? Qtrue : Qfalse;
}

static VALUE udt_rendezvous_set(VALUE klass, VALUE val)
{
	if((TYPE(val) != T_TRUE) && (TYPE(val) != T_FALSE))
		rb_raise(eUDTSocket, "argument error");
	udt_set_bool(fds[klass], UDT_RENDEZVOUS, (val == Qtrue));
	return Qnil;
}

static VALUE udt_sndtimeo_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDT_SNDTIMEO));
}

static VALUE udt_sndtimeo_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDT_SNDTIMEO, FIX2INT(val));
	return Qnil;
}

static VALUE udt_rcvtimeo_get(VALUE klass)
{
	return INT2FIX(udt_get_int(fds[klass], UDT_RCVTIMEO));
}

static VALUE udt_rcvtimeo_set(VALUE klass, VALUE val)
{
	if(TYPE(val) != T_FIXNUM)
		rb_raise(eUDTSocket, "argument error");
	udt_set_int(fds[klass], UDT_RCVTIMEO, FIX2INT(val));
	return Qnil;
}

static VALUE udt_reuseaddr_get(VALUE klass)
{
	return udt_get_bool(fds[klass], UDT_REUSEADDR) ? Qtrue : Qfalse;
}

static VALUE udt_reuseaddr_set(VALUE klass, VALUE val)
{
	if((TYPE(val) != T_TRUE) && (TYPE(val) != T_FALSE))
		rb_raise(eUDTSocket, "argument error");
	udt_set_bool(fds[klass], UDT_REUSEADDR, (val == Qtrue));
	return Qnil;
}

static VALUE udt_maxbw_get(VALUE klass)
{
	return udt_get_int64(fds[klass], UDT_MAXBW);
}

static VALUE udt_maxbw_set(VALUE klass, VALUE val)
{
	int64_t value;
	if(TYPE(val) == T_FIXNUM)
		value = FIX2INT(val);
	else if(TYPE(val) == T_BIGNUM)
		value = FIX2INT(rb_big2ll(val));
	else
		rb_raise(eUDTSocket, "argument error");
	udt_set_int64(fds[klass], UDT_MAXBW, value);
	return Qnil;
}

// initialize

extern "C" void Init_udt()
{
	VALUE cUDTSocket;

	eUDTSocket = rb_define_class("UDTError", rb_eStandardError);
		
	cUDTSocket = rb_define_class("UDTSocket", rb_cObject);
	
	rb_define_singleton_method(cUDTSocket, "startup", (VALUE (*)(ANYARGS))udt_s_startup, 0);
	rb_define_singleton_method(cUDTSocket, "cleanup", (VALUE (*)(ANYARGS))udt_s_cleanup, 0);

	rb_define_method(cUDTSocket, "initialize", (VALUE (*)(ANYARGS))udt_socket, 2);
	rb_define_method(cUDTSocket, "connect", (VALUE (*)(ANYARGS))udt_connect, 1);
	rb_define_method(cUDTSocket, "accept", (VALUE (*)(ANYARGS))udt_accept, 0);
	rb_define_method(cUDTSocket, "close", (VALUE (*)(ANYARGS))udt_close, 0);
	rb_define_method(cUDTSocket, "bind", (VALUE (*)(ANYARGS))udt_bind, 1);
	rb_define_method(cUDTSocket, "listen", (VALUE (*)(ANYARGS))udt_listen, 1);
	rb_define_method(cUDTSocket, "send", (VALUE (*)(ANYARGS))udt_send, 1);
	rb_define_method(cUDTSocket, "recv", (VALUE (*)(ANYARGS))udt_recv, 1);
	rb_define_method(cUDTSocket, "select", (VALUE (*)(ANYARGS))udt_select, -1);

	rb_define_method(cUDTSocket, "udt_mss", (VALUE (*)(ANYARGS))udt_mss_get, 0);
	rb_define_method(cUDTSocket, "udt_mss=", (VALUE (*)(ANYARGS))udt_mss_set, 1);
	rb_define_method(cUDTSocket, "udt_sndsyn?", (VALUE (*)(ANYARGS))udt_sndsyn_get, 0);
	rb_define_method(cUDTSocket, "udt_sndsyn=", (VALUE (*)(ANYARGS))udt_sndsyn_set, 1);
	rb_define_method(cUDTSocket, "udt_rcvsyn?", (VALUE (*)(ANYARGS))udt_sndsyn_get, 0);
	rb_define_method(cUDTSocket, "udt_rcvsyn=", (VALUE (*)(ANYARGS))udt_sndsyn_set, 1);
	rb_define_method(cUDTSocket, "udt_fc", (VALUE (*)(ANYARGS))udt_fc_get, 0);
	rb_define_method(cUDTSocket, "udt_fc=", (VALUE (*)(ANYARGS))udt_fc_set, 1);
	rb_define_method(cUDTSocket, "udt_sndbuf", (VALUE (*)(ANYARGS))udt_sndbuf_get, 0);
	rb_define_method(cUDTSocket, "udt_sndbuf=", (VALUE (*)(ANYARGS))udt_sndbuf_set, 1);
	rb_define_method(cUDTSocket, "udt_rcvbuf", (VALUE (*)(ANYARGS))udt_rcvbuf_get, 0);
	rb_define_method(cUDTSocket, "udt_rcvbuf=", (VALUE (*)(ANYARGS))udt_rcvbuf_set, 1);
	rb_define_method(cUDTSocket, "udp_sndbuf", (VALUE (*)(ANYARGS))udp_sndbuf_get, 0);
	rb_define_method(cUDTSocket, "udp_sndbuf=", (VALUE (*)(ANYARGS))udp_sndbuf_set, 1);
	rb_define_method(cUDTSocket, "udp_rcvbuf", (VALUE (*)(ANYARGS))udp_rcvbuf_get, 0);
	rb_define_method(cUDTSocket, "udp_rcvbuf=", (VALUE (*)(ANYARGS))udp_rcvbuf_set, 1);
	rb_define_method(cUDTSocket, "udt_rendezvous?", (VALUE (*)(ANYARGS))udt_rendezvous_get, 0);
	rb_define_method(cUDTSocket, "udt_rendezvous=", (VALUE (*)(ANYARGS))udt_rendezvous_set, 1);
	rb_define_method(cUDTSocket, "udt_sndtimeo", (VALUE (*)(ANYARGS))udt_sndtimeo_get, 0);
	rb_define_method(cUDTSocket, "udt_sndtimeo=", (VALUE (*)(ANYARGS))udt_sndtimeo_set, 1);
	rb_define_method(cUDTSocket, "udt_rcvtimeo", (VALUE (*)(ANYARGS))udt_rcvtimeo_get, 0);
	rb_define_method(cUDTSocket, "udt_rcvtimeo=", (VALUE (*)(ANYARGS))udt_rcvtimeo_set, 1);
	rb_define_method(cUDTSocket, "udt_reuseaddr?", (VALUE (*)(ANYARGS))udt_reuseaddr_get, 0);
	rb_define_method(cUDTSocket, "udt_reuseaddr=", (VALUE (*)(ANYARGS))udt_reuseaddr_set, 1);
	rb_define_method(cUDTSocket, "udt_maxbw", (VALUE (*)(ANYARGS))udt_maxbw_get, 0);
	rb_define_method(cUDTSocket, "udt_maxbw=", (VALUE (*)(ANYARGS))udt_maxbw_set, 1);
}
